package cn.ewan.cloud.tenant.data.router.mq.core;

import cn.ewan.cloud.common.support.Log;
import cn.ewan.cloud.common.web.support.BusinessException;
import cn.ewan.cloud.common.web.support.rocketmq.RocketMQAsyncCallback;
import cn.ewan.cloud.common.web.support.rocketmq.RocketMQMsg;
import cn.ewan.cloud.common.web.support.rocketmq.RocketMQProducer;
import cn.ewan.cloud.tenant.data.router.core.business.context.TenantContext;
import cn.ewan.cloud.tenant.data.router.core.common.constants.TenantConstants;
import cn.ewan.cloud.tenant.data.router.core.common.exception.TenantMsgCode;
import cn.hutool.core.map.MapUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.client.producer.SendResult;

import java.util.Optional;
import java.util.function.Consumer;

/**
 * @author lgh
 * @date 2023/3/1 17:09
 */
public class TenantRocketMQProducer extends RocketMQProducer {


    @Override
    public <T> void syncSendWithCompensation(RocketMQMsg<T> msg) {
        addTenantTag(msg);
        super.syncSendWithCompensation(msg);
    }


    @Override
    public <T> void asyncSend(RocketMQMsg<T> msg) {
        addTenantTag(msg);
        super.asyncSend(msg);
    }


    @Override
    public <T> SendResult syncSend(RocketMQMsg<T> msg) {
        addTenantTag(msg);
        return super.syncSend(msg);
    }


    @Override
    public <T> void asyncSend(RocketMQMsg<T> msg, Consumer<RocketMQAsyncCallback> callbackProcessor) {
        addTenantTag(msg);
        super.asyncSend(msg, callbackProcessor);
    }


    /**
     * 添加多租户tag
     *
     * @param msg
     */
    private void addTenantTag(RocketMQMsg msg) {
        //因自动补偿机制时, 会有请求头
        String msgTenant = null;
        if (msg != null && msg.getHeaders() != null) {
            msgTenant = (String) msg.getHeaders().get(TenantConstants.TENANT_TAG);
        }
        //若租户取不到值, 则使用请求头的
        String tenant = Optional.ofNullable(TenantContext.getTenantNoCheck()).orElse(msgTenant);
        Log.NORMAL.info("rocket mq tenant interceptor threadId={},tenant={},msgTenant={},topic={}", Thread.currentThread().getId(),
                tenant, msgTenant, msg.getTopic());

        if (StringUtils.isBlank(tenant)) {
            throw new BusinessException(TenantMsgCode.TENANT_FLAG_NO_EXIST);
        }
        if (msg != null) {
            msg.setHeaders(MapUtil.of(TenantConstants.TENANT_TAG, tenant));
        }
    }
}
